/***************************************************************************************************
 * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 *modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright notice,
 *this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *notice, this list of conditions and the following disclaimer in the
 *documentation and/or other materials provided with the distribution.
 *     * Neither the name of the NVIDIA CORPORATION nor the names of its
 *contributors may be used to endorse or promote products derived from this
 *software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 *AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 *IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY DIRECT,
 *INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 *DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 *OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TOR (INCLUDING
 *NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
 *EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 **************************************************************************************************/
/*! \file
    \brief Implicit GEMM testbed sizes for Conv2d problem
*/
#pragma once

#include <vector>

#include "../../common/cutlass_unit_test.h"

#include "cutlass/cutlass.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/conv/convolution.h"
#include "cutlass/conv/conv2d_problem_size.h"

#define CUTLASS_CONV_UNIT_TEST_RIGOROUS_SIZE_ENABLED 1

namespace test {
namespace conv {
namespace device {

using Conv2dProblemVector = std::vector<cutlass::conv::Conv2dProblemSize>;

//
// Structures to prune items from Conv2dProblemVector
//
// Specification template for pruning items for convolution problem lists
template <typename T>
struct Specification {
    virtual ~Specification() = default;
    virtual bool is_satisfied(T item) const = 0;
};

// input size  (NHWC) specification
struct InputSizeSpecification
        : Specification<cutlass::conv::Conv2dProblemSize> {
    cutlass::Tensor4DCoord input_size;

    InputSizeSpecification(cutlass::Tensor4DCoord input_size_)
            : input_size(input_size_) {}

    bool is_satisfied(cutlass::conv::Conv2dProblemSize item) const override {
        return ((input_size.n() == item.N) && (input_size.h() == item.H) &&
                (input_size.w() == item.W) && (input_size.c() == item.C));
    }
};

// stride (stride_h, stride_w) specification
struct StrideSpecification : Specification<cutlass::conv::Conv2dProblemSize> {
    cutlass::MatrixCoord stride;

    StrideSpecification(cutlass::MatrixCoord stride_) : stride(stride_) {}

    bool is_satisfied(cutlass::conv::Conv2dProblemSize item) const override {
        return ((stride.row() == item.stride_h) &&
                (stride.column() == item.stride_h));
    }
};

// channel (C,K) specification, must be multiple of minimum channel
struct ChannelDivisibilitySpecification
        : Specification<cutlass::conv::Conv2dProblemSize> {
    int channel_multiple;

    ChannelDivisibilitySpecification(int channel_multiple_)
            : channel_multiple(channel_multiple_) {}

    bool is_satisfied(cutlass::conv::Conv2dProblemSize item) const override {
        return ((item.K % channel_multiple == 0) &&
                (item.C % channel_multiple == 0));
    }
};

//
// Pruning function for items from Conv2dProblemVector based on a Specification
//
inline Conv2dProblemVector prune(
        Conv2dProblemVector const& items,
        Specification<cutlass::conv::Conv2dProblemSize> const& spec) {
    Conv2dProblemVector pruned_list;

    for (auto& p : items)
        if (spec.is_satisfied(p))
            pruned_list.push_back(p);
    return pruned_list;
}

////////////////////////////////////////////////////////////////////////////
/// Structure TestbedConv2dProblemSizes initializes and holds conv default and
/// important network sizes
////////////////////////////////////////////////////////////////////////////
struct TestbedConv2dProblemSizes {
    //
    // Data members
    //
    int minimum_channel_size;

    Conv2dProblemVector conv2d_default_sizes;
    Conv2dProblemVector conv2d_rigorous_sizes;
    Conv2dProblemVector conv2d_resnet50_sizes;
    Conv2dProblemVector conv2d_resnet50_sizes_perf;

    //
    // Methods
    //
    /// Default ctor
    TestbedConv2dProblemSizes(int minimum_channel_size_ = 64)
            : minimum_channel_size(minimum_channel_size_) {
        initialize_conv2d_default_sizes();
        initialize_conv2d_rigorous_sizes();
        initialize_conv2d_resnet50_sizes(conv2d_resnet50_sizes,
                                         1 /*batch-size*/);

        initialize_conv2d_resnet50_sizes(conv2d_resnet50_sizes_perf,
                                         34 /*batch-size*/);
        filter_all();
    }

    /// Eliminates some illegal cases
    void filter_all() {
        Conv2dProblemVector* problems_vectors[] = {
                &conv2d_default_sizes, &conv2d_rigorous_sizes,
                &conv2d_resnet50_sizes, &conv2d_resnet50_sizes_perf};

        for (Conv2dProblemVector* problems : problems_vectors) {
            Conv2dProblemVector filtered;

            for (cutlass::conv::Conv2dProblemSize const& problem : *problems) {
                if ((!(problem.C % minimum_channel_size)) &&
                    (!(problem.K % minimum_channel_size))) {
                    filtered.push_back(problem);
                }
            }

            *problems = filtered;
        }
    }

    // Add a few standard convolution problem sizes
    void initialize_conv2d_default_sizes() {
        ////////////////////////////////////////////////////////////////////////////////////////////
        // Very Small input size (1x8x8xminimum_channel_size), filter size (3x3
        // - 7x7), stride (1,1) C < CTA::K and non-multiples of CTA::K. Typical
        // CTA::K = {32, 64}
        ////////////////////////////////////////////////////////////////////////////////////////////

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 8, 8, minimum_channel_size},  // input size  (NHWC)
                {minimum_channel_size, 3, 3,
                 minimum_channel_size},  // filter size (KRSC)
                {1, 1, 1, 1},            // padding (pad_h, _, pad_w, _)
                {1, 1},                  // stride (stride_h, stride_w)
                {1, 1}                   // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 8, 8, minimum_channel_size},  // input size  (NHWC)
                {minimum_channel_size, 4, 4,
                 minimum_channel_size},  // filter size (KRSC)
                {1, 1, 1, 1},            // padding (pad_h, _, pad_w, _)
                {1, 1},                  // stride (stride_h, stride_w)
                {1, 1}                   // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 8, 8, minimum_channel_size},  // input size  (NHWC)
                {minimum_channel_size, 5, 5,
                 minimum_channel_size},  // filter size (KRSC)
                {1, 1, 1, 1},            // padding (pad_h, _, pad_w, _)
                {1, 1},                  // stride (stride_h, stride_w)
                {1, 1}                   // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 8, 8, minimum_channel_size},  // input size  (NHWC)
                {minimum_channel_size, 6, 5,
                 minimum_channel_size},  // filter size (KRSC)
                {1, 1, 1, 1},            // padding (pad_h, _, pad_w, _)
                {1, 1},                  // stride (stride_h, stride_w)
                {1, 1}                   // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 8, 8, minimum_channel_size},  // input size  (NHWC)
                {minimum_channel_size, 6, 6,
                 minimum_channel_size},  // filter size (KRSC)
                {1, 1, 1, 1},            // padding (pad_h, _, pad_w, _)
                {1, 1},                  // stride (stride_h, stride_w)
                {1, 1}                   // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 8, 8, minimum_channel_size},  // input size  (NHWC)
                {minimum_channel_size, 7, 7,
                 minimum_channel_size},  // filter size (KRSC)
                {1, 1, 1, 1},            // padding (pad_h, _, pad_w, _)
                {1, 1},                  // stride (stride_h, stride_w)
                {1, 1}                   // dilation (dilation_h, dilation_w)
                ));

        ////////////////////////////////////////////////////////////////////////////////////
        // Medium input size (1x16x16x128), filter size (1x1, 2x2, 3x3, 5x5),
        // stride (1, 1)
        ////////////////////////////////////////////////////////////////////////////////////
        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 15, 19, 160},  // input size  (NHWC)
                {224, 1, 1, 160},  // filter size (KRSC)
                {0, 0, 0, 0},      // padding (pad_h, _, pad_w, _)
                {1, 1},            // stride (stride_h, stride_w)
                {1, 1}             // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 16, 16, 160},  // input size  (NHWC)
                {224, 2, 3, 160},  // filter size (KRSC)
                {1, 1, 1, 1},      // padding (pad_h, _, pad_w, _)
                {1, 1},            // stride (stride_h, stride_w)
                {1, 1}             // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 23, 21, 128},  // input size  (NHWC)
                {224, 3, 3, 128},  // filter size (KRSC)
                {1, 1, 1, 1},      // padding (pad_h, _, pad_w, _)
                {1, 1},            // stride (stride_h, stride_w)
                {1, 1}             // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 29, 37, 160},  // input size  (NHWC)
                {224, 5, 5, 160},  // filter size (KRSC)
                {2, 2, 2, 2},      // padding (pad_h, _, pad_w, _)
                {1, 1},            // stride (stride_h, stride_w)
                {1, 1}             // dilation (dilation_h, dilation_w)
                ));

        ////////////////////////////////////////////////////////////////////////////////////
        // C > CTA::K and non-multiples of CTA::K. Typical CTA::K = {32, 64}
        ////////////////////////////////////////////////////////////////////////////////////
        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 15, 19, 32 + minimum_channel_size},  // input size  (NHWC)
                {96, 3, 3, 32 + minimum_channel_size},   // filter size (KRSC)
                {1, 1, 1, 1},  // padding (pad_h, _, pad_w, _)
                {1, 1},        // stride (stride_h, stride_w)
                {1, 1}         // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 16, 16, 64 + minimum_channel_size},  // input size  (NHWC)
                {96, 3, 3, 64 + minimum_channel_size},   // filter size (KRSC)
                {1, 1, 1, 1},  // padding (pad_h, _, pad_w, _)
                {1, 1},        // stride (stride_h, stride_w)
                {1, 1}         // dilation (dilation_h, dilation_w)
                ));

        ////////////////////////////////////////////////////////////////////////////////////
        // Medium input size (1x16x16x128), filter size (1x1, 3,x3, 5x5), stride
        // (2, 2)
        ////////////////////////////////////////////////////////////////////////////////////
        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 19, 37, 160},  // input size  (NHWC)
                {224, 3, 3, 160},  // filter size (KRSC)
                {1, 1, 1, 1},      // padding (pad_h, _, pad_w, _)
                {2, 2},            // stride (stride_h, stride_w)
                {1, 1}             // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 16, 16, 288},  // input size  (NHWC)
                {160, 5, 5, 288},  // filter size (KRSC)
                {2, 2, 2, 2},      // padding (pad_h, _, pad_w, _)
                {2, 2},            // stride (stride_h, stride_w)
                {1, 1}             // dilation (dilation_h, dilation_w)
                ));

        /////////////////////////////////////////////////////////////////////////////
        // Additional input size
        /////////////////////////////////////////////////////////////////////////////
        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {3, 28, 28, 256},  // input size  (NHWC)
                {256, 2, 2, 256},  // filter size (KRSC)
                {0, 0, 0, 0},      // padding (pad_h, _, pad_w, _)
                {2, 2},            // stride (stride_h, stride_w)
                {1, 1}             // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {32, 32, 32, 32},  // input size  (NHWC)
                {32, 1, 1, 32},    // filter size (KRSC)
                {0, 0, 0, 0},      // padding (pad_h, _, pad_w, _)
                {1, 1},            // stride (stride_h, stride_w)
                {1, 1}             // dilation (dilation_h, dilation_w)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {4, 3, 3, 128},    // input size  (NHWC)
                {256, 3, 3, 128},  // filter size (KRSC)
                {0, 0, 0, 0},      // padding (pad_h, _, pad_w, _)
                {1, 1},            // stride (stride_h, stride_w)
                {1, 1},            // dilation (dilation_h, dilation_w)
                {4, 3, 3, 256}     // output size (NPQK)
                ));

        conv2d_default_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {4, 1, 1, 256},    // input size  (NHWC)
                {328, 3, 3, 256},  // filter size (KRSC)
                {1, 1, 1, 1},      // padding (pad_h, _, pad_w, _)
                {1, 1},            // stride (stride_h, stride_w)
                {1, 1},            // dilation (dilation_h, dilation_w)
                {4, 1, 1, 328}     // output size (NPQK)
                ));
    }

    // Add a few large and rigorous convolution problem sizes
    void initialize_conv2d_rigorous_sizes() {
#if CUTLASS_CONV_UNIT_TEST_RIGOROUS_SIZE_ENABLED
        conv2d_rigorous_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 124, 224, 96},  // input size  (NHWC)
                {24, 7, 7, 96},     // filter size (KRSC)
                {1, 229, 129, 32}   // output size (NPQK)
                ));

        conv2d_rigorous_sizes.push_back(cutlass::conv::Conv2dProblemSize(
                {1, 233, 35, 48},  // input size  (NHWC)
                {24, 7, 5, 48},    // filter size (KRSC)
                {1, 233, 35, 24}   // output size (NPQK)
                ));

#endif
    }

    // Add resent50 layers to unit testing sizes
    void initialize_conv2d_resnet50_sizes(
            Conv2dProblemVector& conv2d_problem_vector, int batch_size = 1) {
#if 0  // Resnet50 first layer (layer_id = 0) with channel = 3 is not supported
       // in cutlass
    conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(   
      [1, 224, 224, 3],           // input size (NHWC)
      [64, 7, 7, 3],              // filter size (KRSC)
      [3, 3, 3, 3],               // padding (pad_h, _, pad_w, _)
      [2, 2],                     // stride (stride_h, stride_w)
      [1, 1],                     // dilation (dilation_h, dilation_w)
    ));
#endif

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 56, 56, 64},  // input size (NHWC)
                {256, 1, 1, 64},           // filter size (KRSC)
                {0, 0, 0, 0},              // padding (pad_h, _, pad_w, _)
                {1, 1},                    // stride (stride_h, stride_w)
                {1, 1}                     // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 56, 56, 64},  // input size (NHWC)
                {64, 1, 1, 64},            // filter size (KRSC)
                {0, 0, 0, 0},              // padding (pad_h, _, pad_w, _)
                {1, 1},                    // stride (stride_h, stride_w)
                {1, 1}                     // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 56, 56, 64},  // input size (NHWC)
                {64, 3, 3, 64},            // filter size (KRSC)
                {1, 1, 1, 1},              // padding (pad_h, _, pad_w, _)
                {1, 1},                    // stride (stride_h, stride_w)
                {1, 1}                     // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 56, 56, 256},  // input size (NHWC)
                {64, 1, 1, 256},            // filter size (KRSC)
                {0, 0, 0, 0},               // padding (pad_h, _, pad_w, _)
                {1, 1},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 56, 56, 256},  // input size (NHWC)
                {512, 1, 1, 256},           // filter size (KRSC)
                {0, 0, 0, 0},               // padding (pad_h, _, pad_w, _)
                {2, 2},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 56, 56, 256},  // input size (NHWC)
                {128, 1, 1, 256},           // filter size (KRSC)
                {0, 0, 0, 0},               // padding (pad_h, _, pad_w, _)
                {2, 2},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 28, 28, 128},  // input size (NHWC)
                {128, 3, 3, 128},           // filter size (KRSC)
                {1, 1, 1, 1},               // padding (pad_h, _, pad_w, _)
                {1, 1},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 28, 28, 128},  // input size (NHWC)
                {512, 1, 1, 128},           // filter size (KRSC)
                {0, 0, 0, 0},               // padding (pad_h, _, pad_w, _)
                {1, 1},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 28, 28, 512},  // input size (NHWC)
                {128, 1, 1, 512},           // filter size (KRSC)
                {0, 0, 0, 0},               // padding (pad_h, _, pad_w, _)
                {1, 1},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 28, 28, 512},  // input size (NHWC)
                {1024, 1, 1, 512},          // filter size (KRSC)
                {0, 0, 0, 0},               // padding (pad_h, _, pad_w, _)
                {2, 2},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 28, 28, 512},  // input size (NHWC)
                {256, 1, 1, 512},           // filter size (KRSC)
                {0, 0, 0, 0},               // padding (pad_h, _, pad_w, _)
                {2, 2},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 14, 14, 256},  // input size (NHWC)
                {256, 3, 3, 256},           // filter size (KRSC)
                {1, 1, 1, 1},               // padding (pad_h, _, pad_w, _)
                {1, 1},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 14, 14, 256},  // input size (NHWC)
                {1024, 1, 1, 256},          // filter size (KRSC)
                {0, 0, 0, 0},               // padding (pad_h, _, pad_w, _)
                {1, 1},                     // stride (stride_h, stride_w)
                {1, 1}                      // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 14, 14, 1024},  // input size (NHWC)
                {256, 1, 1, 1024},           // filter size (KRSC)
                {0, 0, 0, 0},                // padding (pad_h, _, pad_w, _)
                {1, 1},                      // stride (stride_h, stride_w)
                {1, 1}  // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 14, 14, 1024},  // input size (NHWC)
                {2048, 1, 1, 1024},          // filter size (KRSC)
                {0, 0, 0, 0},                // padding (pad_h, _, pad_w, _)
                {2, 2},                      // stride (stride_h, stride_w)
                {1, 1}  // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 14, 14, 1024},  // input size (NHWC)
                {512, 1, 1, 1024},           // filter size (KRSC)
                {0, 0, 0, 0},                // padding (pad_h, _, pad_w, _)
                {2, 2},                      // stride (stride_h, stride_w)
                {1, 1}  // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 7, 7, 512},  // input size (NHWC)
                {512, 3, 3, 512},         // filter size (KRSC)
                {1, 1, 1, 1},             // padding (pad_h, _, pad_w, _)
                {1, 1},                   // stride (stride_h, stride_w)
                {1, 1}                    // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 7, 7, 512},  // input size (NHWC)
                {2048, 1, 1, 512},        // filter size (KRSC)
                {0, 0, 0, 0},             // padding (pad_h, _, pad_w, _)
                {1, 1},                   // stride (stride_h, stride_w)
                {1, 1}                    // dilation (dilation_h, dilation_w)
                ));

        conv2d_problem_vector.push_back(cutlass::conv::Conv2dProblemSize(
                {batch_size, 7, 7, 2048},  // input size (NHWC)
                {512, 1, 1, 2048},         // filter size (KRSC)
                {0, 0, 0, 0},              // padding (pad_h, _, pad_w, _)
                {1, 1},                    // stride (stride_h, stride_w)
                {1, 1}                     // dilation (dilation_h, dilation_w)
                ));
    }
};

}  // namespace device
}  // namespace conv
}  // namespace test
